from web_apps.llm.llm_utils import extract_code, process_dataframe
import traceback
from utils.common_utils import get_now_time


class DataChatAgent:
    def __init__(self, llm, reader, knowledge='', answer='', retry=1, max_token=4000):
        self.llm = llm
        self.reader = reader
        self.knowledge = knowledge
        self.answer = answer
        self.question = ''
        self.last_code_executed = ''
        self.code_exception = ''
        self.max_retry = retry
        self.info_prompt = ''
        self.question = ''
        self.max_token = max_token

    def gen_info_prompt(self):
        '''
        生成信息提示
        :return:
        '''
        if self.info_prompt == '':
            info_prompt = self.reader.get_info_prompt('')
            if len(info_prompt) > self.max_token:
                # 信息过长，抽出模型列表，使用llm筛选出部分模型生成信息提示
                model_list = self.reader.gen_models()
                model_list = [{'type': i['type'], 'name': i['model_conf']['name']} for i in model_list]
                prompt = f"你正在进行数据分析任务，有以下数据模型：\n{model_list}\n 请根据问题：\n {self.question}\n 从以上数据模型中筛选出你需要的模型名称列表,只需要返回名称列表，用逗号隔开，不要其他内容"
                model_prompt = self.llm.invoke(prompt).content
                info_prompt = self.reader.get_info_prompt(model_prompt)
            self.info_prompt = info_prompt
        return self.info_prompt

    def generate_code(self, prompt):
        result_example_prompt = '{ "type": "string", "value": "100" } or { "type": "dataframe", "value": pd.DataFrame({...}) } or { "type": "html", "value": line.render_embed() }'
        prompt = f"""
I have a data reader object called reader, and the object information is：
{self.gen_info_prompt()}

Update this initial code:
```python
# TODO: import the required dependencies

# Write code here

# Declare result var: 
type (possible values "string", "dataframe", "html"). Example: {result_example_prompt}

```

### QUERY

{prompt}

Variable `reader` is already declared.

At the end, declare "result" variable as a dictionary of type and value.

If you are asked to plot a chart, use "pyecharts" for charts, use the render_embed() function to return the corresponding html type and the html content value.

Generate python code and return full updated code:
请在代码中使用中文添加必要注释
"""
        if self.knowledge != '':
            prompt = f"结合知识库信息:\n{self.knowledge}\n回答以下问题:\n{prompt}"
        self.llm_result = self.llm.invoke(prompt).content
        code = extract_code(self.llm_result)
        return code

    def fix_code(self):
        '''
        使用llm修正错误代码
        :return:
        '''
        fix_code_prompt = f"""
I have a data reader object called reader, and the object information is：
{self.gen_info_prompt()}
The user asked the following question:
{self.question}
You generated this python code:
{self.last_code_executed}
the code running throws an exception:
{self.code_exception}
Fix the python code above and return the new python code
请在代码中使用中文添加必要注释
        """
        if self.knowledge != '':
            fix_code_prompt = f"结合知识库信息:\n{self.knowledge}\n回答以下问题:\n{fix_code_prompt}"
        self.llm_result = self.llm.invoke(fix_code_prompt).content
        new_code = extract_code(self.llm_result)
        return new_code

    def execute_code(self, code: str):
        """
        Execute the python code generated by LLMs to answer the question
        about the input dataframe. Run the code in the current context and return the
        result.
        Args:
            code (str): Python code to execute.
            context (CodeExecutionContext): Code Execution Context
                    with prompt id and skills.
        Returns:
            Any: The result of the code execution. The type of the result depends
                on the generated code.
        """
        try:
            environment = {'reader': self.reader}
            exec(code, environment)
            self.last_code_executed = code
            if "result" not in environment:
                raise ValueError("No result returned")
            else:
                result = environment['result']
                return result
        except Exception as e:
            self.last_code_executed = code
            raise e

    def parse_result(self, result):
        if result['type'] == 'html':
            return {'content': result['value'], 'type': 'html'}
        elif result['type'] == 'dataframe':
            data_li = process_dataframe(result)
            return {'content': data_li, 'type': 'data'}
        else:
            return {'content': result['value'], 'type': 'text'}

    def run(self, prompt):
        self.question = prompt
        retry_count = 0
        result = None
        code = ''
        while retry_count < self.max_retry:
            try:
                if retry_count == 0:
                    if self.answer != '':
                        self.gen_info_prompt()  # 执行一次部分数据模型初始化
                        self.llm_result = self.answer
                        code = extract_code(self.answer)
                    else:
                        code = self.generate_code(prompt)
                result = self.execute_code(code)
                result = self.parse_result(result)
                return result
            except Exception as e:
                traceback_errors = traceback.format_exc()
                self.code_exception = traceback_errors
                retry_count += 1
                code = self.fix_code()
        return result

    def chat(self, prompt):
        self.question = prompt
        retry_count = 0
        result = None
        code = ''
        while retry_count <= self.max_retry:
            try:
                if retry_count == 0:
                    if self.answer != '':
                        data = {'content': {'title': '生成处理代码', 'content': '发现知识库中答案，直接使用',
                                            'time': get_now_time(res_type='datetime')}, 'type': 'flow'}
                        yield data
                        self.gen_info_prompt()  # 执行一次部分数据模型初始化
                        self.llm_result = self.answer
                        code = extract_code(self.answer)
                    else:
                        data = {'content': {'title': '生成处理代码', 'content': '开始生成处理代码',
                                            'time': get_now_time(res_type='datetime')}, 'type': 'flow'}
                        yield data
                        code = self.generate_code(prompt)
                    data = {'content': {'title': '处理代码生成成功', 'content': self.llm_result,
                                        'time': get_now_time(res_type='datetime')}, 'type': 'flow'}
                    yield data
                data = {'content': {'title': '执行处理代码', 'content': f"```python\n{code}\n```", 'time': get_now_time(res_type='datetime')}, 'type': 'flow'}
                yield data
                result = self.execute_code(code)
                data = {'content': {'title': '处理完成', 'content': '处理完成', 'time': get_now_time(res_type='datetime')}, 'type': 'flow'}
                yield data
                data = self.parse_result(result)
                yield data
                return result
            except Exception as e:
                traceback_errors = traceback.format_exc()
                self.code_exception = traceback_errors
                data = {'content': {'title': '执行代码出错，修复代码', 'content': f'执行代码报错：{self.code_exception}', 'time': get_now_time(res_type='datetime')}, 'type': 'flow'}
                yield data
                retry_count += 1
                code = self.fix_code()
                data = {'content': {'title': '修复处理代码成功', 'content': self.llm_result, 'time': get_now_time(res_type='datetime')}, 'type': 'flow'}
                yield data
        data = {'content': {'title': '处理失败', 'content': f'处理失败：{self.code_exception}', 'time': get_now_time(res_type='datetime')}, 'type': 'flow'}
        yield data
        return result
